package com.yupi.yuaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.yupi.yuaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;



/**
 *
 */
@Slf4j
@Data
@EqualsAndHashCode(callSuper = true)
public class ToolCallAgent extends ReActAgent {

    //可用的工具列表
    private final ToolCallback[] availableTools;

    //保存工具调用的聊天响应
    private ChatResponse toolCallChatResponse;

    //工具调用管理器
    private final ToolCallingManager toolCallingManager;

    //禁用内置工具调用，自己实现工具调用
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true) // 禁用Spring AI内置的工具调用，由我们自主实现
                .build();
    }

    @Override
    public boolean think(){
        if (getNextStepPrompt() != null && !getNextStepPrompt().isEmpty()) {
            UserMessage userMessage = new UserMessage(getNextStepPrompt());
            getMessageList().add(userMessage);
        }
        List<Message> messageList = getMessageList();
        Prompt prompt = new Prompt(messageList, chatOptions);
        try {
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            this.toolCallChatResponse = chatResponse;
            AssistantMessage assistantMessage = null;
            if (chatResponse != null) {
                assistantMessage = chatResponse.getResult().getOutput();
            }

            String result = null;
            if (assistantMessage != null) {
                result = assistantMessage.getText();
            }
            log.info("{}的思考: {}", getName(), result);

            List<AssistantMessage.ToolCall> toolCallList = null;
            if (assistantMessage != null) {
                toolCallList = assistantMessage.getToolCalls();
                log.info("{}选择了 {} 个工具调用", getName(),  toolCallList.size());
            }
            String toolCallInfo = toolCallList.stream()
                    .map(toolCall -> String.format("工具名称：%s，参数：%s", toolCall.name(), toolCall.arguments()))
                    .collect(Collectors.joining("\n"));
            log.info("工具调用详情：\n{}", toolCallInfo);
            if (toolCallList.isEmpty()) {
                getMessageList().add(assistantMessage);
                return false;
            }else {
                // 调用工具时，不需要深拷贝上下文，因为工具调用时，上下文已经记录在toolCallChatResponse中
                return true;
            }
        } catch (Exception e) {
            log.error("{} 思考过程遇到了问题: {}", getName(), e.getMessage());
            getMessageList().add(new AssistantMessage("抱歉，我的思考过程出现了问题，请稍后再试" + e.getMessage()));
            return false;
        }
    }




    @Override
    public String act() {
        if (toolCallChatResponse.hasToolCalls()) {
            log.info("有工具调用");
        }
        // 工具调用
        Prompt prompt = new Prompt(getMessageList(), chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        // 这里的 conversationHistory 已经包含了助手承接工具调用回应结果
        setMessageList(toolExecutionResult.conversationHistory());
        // 抽取工具调用的返回信息
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        String results = toolResponseMessage.getResponses().stream()
                .map(response -> "工具：" + response.name() + " 完成了你的任务！结果：" + response.responseData())
                .collect(Collectors.joining("\n"));
        // 判断是否调用了终止工具
        if (toolResponseMessage.getResponses().stream().anyMatch(toolResponse -> toolResponse.name().equals("doTerminate"))) {
            setState(AgentState.FINISHED);
        }

        log.info("工具调用返回结果：\n{}", results);
        return results;
    }


    @Override
    public String step() {
        return super.step();
    }
}
